/**
 * Main ComfyUI client class
 */

import type {
  ComfyUIClientConfig,
  ExecutionCallbacks,
  ExecutionOptions,
  TemplateExecutionResult,
  ParameterValues,
  ComfyUIWorkflow
} from '../types';
import { HTTPClient } from './HTTPClient';
import { WebSocketClient } from './WebSocketClient';
import { WorkflowTemplate } from '../templates/WorkflowTemplate';
import { TemplateManager } from '../templates/TemplateManager';
import { 
  ComfyUIError, 
  TimeoutError, 
  ExecutionInterruptedError,
  withTimeout 
} from '../utils/error-handling';

/**
 * Execution context for tracking ongoing executions
 */
interface ExecutionContext {
  promptId: string;
  templateId: string | undefined;
  callbacks: ExecutionCallbacks | undefined;
  startTime: number;
  timeout: number | undefined;
  outputs: Record<string, unknown>;
  resolve: (result: TemplateExecutionResult) => void;
  reject: (error: Error) => void;
}

/**
 * Main ComfyUI client for workflow execution and management
 */
export class ComfyUIClient {
  private readonly httpClient: HTTPClient;
  private readonly wsClient: WebSocketClient;
  private readonly templateManager: TemplateManager;
  private readonly executions = new Map<string, ExecutionContext>();
  private isConnected = false;

  /**
   * Creates a new ComfyUIClient instance
   * @param config - Client configuration
   */
  constructor(config: ComfyUIClientConfig) {
    this.httpClient = new HTTPClient(config);
    
    // Create WebSocket URL from HTTP URL
    const wsUrl = config.baseUrl
      .replace(/^http/, 'ws')
      .replace(/\/$/, '') + '/ws';
    
    this.wsClient = new WebSocketClient({
      url: wsUrl,
      reconnectAttempts: 5,
      reconnectDelay: 1000
    });

    this.templateManager = new TemplateManager();
    this.setupEventListeners();
  }

  /**
   * Gets the template manager
   */
  get templates(): TemplateManager {
    return this.templateManager;
  }

  /**
   * Gets the HTTP client
   */
  get http(): HTTPClient {
    return this.httpClient;
  }

  /**
   * Gets the WebSocket client
   */
  get websocket(): WebSocketClient {
    return this.wsClient;
  }

  /**
   * Connects to the ComfyUI server
   * @returns Promise that resolves when connected
   */
  async connect(): Promise<void> {
    if (this.isConnected) {
      return;
    }

    try {
      await this.wsClient.connect();
      this.isConnected = true;
    } catch (error) {
      throw new ComfyUIError(
        'Failed to connect to ComfyUI server',
        'CONNECTION_ERROR',
        error
      );
    }
  }

  /**
   * Disconnects from the ComfyUI server
   */
  disconnect(): void {
    this.wsClient.disconnect();
    this.isConnected = false;
    
    // Reject all pending executions
    for (const [promptId, context] of this.executions) {
      context.reject(new ExecutionInterruptedError('Client disconnected'));
      this.executions.delete(promptId);
    }
  }

  /**
   * Executes a workflow template with parameters
   * @param template - Workflow template to execute
   * @param parameters - Parameter values
   * @param options - Execution options
   * @returns Promise resolving to execution result
   */
  async executeTemplate(
    template: WorkflowTemplate,
    parameters: ParameterValues,
    options: ExecutionOptions = {}
  ): Promise<TemplateExecutionResult> {
    // Ensure connection
    if (!this.isConnected) {
      await this.connect();
    }

    // Create workflow instance
    const instance = template.createInstance(parameters);
    const workflow = instance.getExecutableWorkflow();

    return this.executeWorkflow(workflow, {
      ...options,
      templateId: template.id
    });
  }

  /**
   * Executes a template by ID
   * @param templateId - Template ID to execute
   * @param parameters - Parameter values
   * @param options - Execution options
   * @returns Promise resolving to execution result
   */
  async executeTemplateById(
    templateId: string,
    parameters: ParameterValues,
    options: ExecutionOptions = {}
  ): Promise<TemplateExecutionResult> {
    const template = this.templateManager.getById(templateId);
    return this.executeTemplate(template, parameters, options);
  }

  /**
   * Executes a raw workflow
   * @param workflow - Workflow to execute
   * @param options - Execution options
   * @returns Promise resolving to execution result
   */
  async executeWorkflow(
    workflow: ComfyUIWorkflow,
    options: ExecutionOptions & { templateId?: string } = {}
  ): Promise<TemplateExecutionResult> {
    const startTime = Date.now();

    try {
      // Queue the prompt
      const queueResponse = await this.httpClient.queuePrompt({
        prompt: workflow,
        client_id: this.wsClient.clientId
      });

      // Create execution context
      const executionPromise = new Promise<TemplateExecutionResult>((resolve, reject) => {
        const context: ExecutionContext = {
          promptId: queueResponse.prompt_id,
          templateId: options.templateId,
          callbacks: options.callbacks,
          startTime,
          timeout: options.timeout,
          outputs: {},
          resolve,
          reject
        };

        this.executions.set(queueResponse.prompt_id, context);

        // Set up timeout if specified
        if (options.timeout) {
          setTimeout(() => {
            if (this.executions.has(queueResponse.prompt_id)) {
              this.executions.delete(queueResponse.prompt_id);
              reject(new TimeoutError(`Execution timed out after ${options.timeout}ms`));
            }
          }, options.timeout);
        }
      });

      // Apply timeout wrapper if specified
      if (options.timeout) {
        return await withTimeout(
          executionPromise,
          options.timeout,
          'Template execution timed out'
        );
      }

      return await executionPromise;

    } catch (error) {
      const executionTime = Date.now() - startTime;
      
      return {
        promptId: '',
        success: false,
        error: {
          message: error instanceof Error ? error.message : String(error),
          timestamp: new Date()
        },
        executionTime
      };
    }
  }

  /**
   * Interrupts the current execution
   * @returns Promise that resolves when interrupt is sent
   */
  async interrupt(): Promise<void> {
    await this.httpClient.interrupt();
    
    // Mark all executions as interrupted
    for (const [promptId, context] of this.executions) {
      context.reject(new ExecutionInterruptedError('Execution interrupted'));
      this.executions.delete(promptId);
    }
  }

  /**
   * Clears the execution queue
   * @returns Promise that resolves when queue is cleared
   */
  async clearQueue(): Promise<void> {
    await this.httpClient.clearQueue();
  }

  /**
   * Gets the current queue status
   * @returns Promise resolving to queue status
   */
  async getQueueStatus() {
    return this.httpClient.getQueue();
  }

  /**
   * Gets execution history
   * @param promptId - Optional specific prompt ID
   * @returns Promise resolving to history
   */
  async getHistory(promptId?: string) {
    return this.httpClient.getHistory(promptId);
  }

  /**
   * Gets system statistics
   * @returns Promise resolving to system stats
   */
  async getSystemStats() {
    return this.httpClient.getSystemStats();
  }

  /**
   * Gets object information
   * @returns Promise resolving to object info
   */
  async getObjectInfo() {
    return this.httpClient.getObjectInfo();
  }

  /**
   * Converts execution outputs to HTTP URLs
   * @param outputs - Execution outputs from ComfyUI
   * @returns Array of image URLs or single URL if only one image
   */
  outputsToUrls(outputs: Record<string, unknown>): string | string[] {
    const urls: string[] = [];

    // Iterate through all output nodes
    for (const nodeOutput of Object.values(outputs)) {
      if (nodeOutput && typeof nodeOutput === 'object' && 'images' in nodeOutput) {
        const images = (nodeOutput as any).images;
        if (Array.isArray(images)) {
          for (const image of images) {
            if (image && typeof image === 'object' && 'filename' in image) {
              const url = this.httpClient.getImageUrl(
                image.filename,
                image.subfolder || '',
                image.type || 'output'
              );
              urls.push(url);
            }
          }
        }
      }
    }

    // Return single URL if only one image, otherwise return array
    if (urls.length === 0) {
      return [];
    }
    return urls.length === 1 ? urls[0]! : urls;
  }

  /**
   * Uploads an image from local file path to ComfyUI server
   * @param filePath - Local file path (e.g., "C:\\Users\\user\\image.png")
   * @param subfolder - Optional subfolder on server
   * @param overwrite - Whether to overwrite existing file
   * @returns Promise resolving to server filename
   */
  async uploadImageFromPath(
    filePath: string,
    subfolder?: string,
    overwrite: boolean = false
  ): Promise<string> {
    const result = await this.httpClient.uploadImageFromPath(filePath, subfolder, overwrite);
    return result.name;
  }

  /**
   * Uploads an image and executes a template in one step
   * @param template - Template to execute
   * @param parameters - Template parameters (input_image will be replaced with uploaded filename)
   * @param localImagePath - Local path to image file
   * @param options - Execution options
   * @returns Promise resolving to execution result
   */
  async executeTemplateWithLocalImage(
    template: WorkflowTemplate,
    parameters: ParameterValues,
    localImagePath: string,
    options: ExecutionOptions = {}
  ): Promise<TemplateExecutionResult> {
    // Upload the image first
    console.log(`📤 Uploading image from: ${localImagePath}`);
    const uploadedFilename = await this.uploadImageFromPath(localImagePath);
    console.log(`✅ Image uploaded as: ${uploadedFilename}`);

    // Replace input_image parameter with uploaded filename
    const updatedParameters = {
      ...parameters,
      input_image: uploadedFilename
    };

    // Execute the template with updated parameters
    return this.executeTemplate(template, updatedParameters, options);
  }

  /**
   * Sets up WebSocket event listeners
   * @private
   */
  private setupEventListeners(): void {
    this.wsClient.on('progress', (event) => {
      if (event.type === 'progress') {
        const context = this.executions.get(event.promptId || '');
        if (context?.callbacks?.onProgress) {
          context.callbacks.onProgress({
            nodeId: event.nodeId,
            progress: event.progress,
            max: event.max,
            timestamp: event.timestamp
          });
        }
      }
    });

    this.wsClient.on('executing', (event) => {
      if (event.type === 'executing') {
        const context = this.executions.get(event.promptId || '');
        if (context?.callbacks?.onExecuting) {
          context.callbacks.onExecuting(event.nodeId);
        }
      }
    });

    this.wsClient.on('executed', (event) => {
      if (event.type === 'executed') {
        const context = this.executions.get(event.promptId || '');
        if (context) {
          // Accumulate outputs from each executed node
          if (event.nodeId && event.outputs) {
            context.outputs[event.nodeId] = event.outputs;
          }

          // Call the callback if provided
          if (context.callbacks?.onExecuted) {
            context.callbacks.onExecuted({
              promptId: event.promptId || '',
              outputs: event.outputs,
              executionTime: event.executionTime,
              timestamp: event.timestamp
            });
          }
        }
      }
    });

    this.wsClient.on('completed', (event) => {
      if (event.type === 'completed') {
        const context = this.executions.get(event.promptId || '');
        if (context) {
          const executionTime = Date.now() - context.startTime;

          // Use the accumulated outputs from executed events for immediate response
          const outputs = context.outputs;

          context.resolve({
            promptId: event.promptId || '',
            success: true,
            outputs,
            executionTime
          });
          this.executions.delete(event.promptId || '');
        }
      }
    });

    this.wsClient.on('error', (event) => {
      if (event.type === 'error') {
        const context = this.executions.get(event.promptId || '');
        if (context) {
          if (context.callbacks?.onError) {
            const errorObj: any = {
              message: event.error.message,
              details: event.details,
              timestamp: event.timestamp
            };
            if (event.nodeId) {
              errorObj.nodeId = event.nodeId;
            }
            context.callbacks.onError(errorObj);
          }

          const executionTime = Date.now() - context.startTime;
          context.resolve({
            promptId: event.promptId || '',
            success: false,
            error: {
              message: event.error.message,
              timestamp: event.timestamp
            },
            executionTime
          });
          this.executions.delete(event.promptId || '');
        }
      }
    });

    this.wsClient.on('disconnected', () => {
      this.isConnected = false;
    });

    this.wsClient.on('connected', () => {
      this.isConnected = true;
    });
  }
}
